#!/usr/bin/python
# -*- coding:utf-8 -*-
# @FileName : DL5_test4_1.py
# Author    : myh

import torch
import torch.nn.functional as F
from torch import nn


class MyLinear1(nn.Module):
    def __init__(self, A , B ):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(A, B))
        print(self.weight.data)
        print(self.weight.data[1,4])
    def forward(self, X):
        y = 0
        c,k = X.shape
        for i in range(c):
            x1 = X[:,i]
            for j in range(c):
                x2 = X[:,j]
                y_1 = torch.mul(x1,x2)
                y = y + torch.mul(y_1, self.weight.data[i,j])
        return y

# linear = MyLinear1(5, 5)
# X = torch.randn(2,5)
# print(X)
# print(linear(X))


class MyLinear2(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, X):
        FFT = torch.fft.fft(X)
        return FFT[:, :round(X.shape[1] / 2)]


layer = MyLinear2()
print(layer(torch.rand(3, 7)))


